import torch

device = (torch.device('cuda') if torch.cuda.is_available()
          else torch.device("cpu"))


def validate(model, train_loader, val_loader):
    accdict = {}
    for name, loader in [("train", train_loader), ("val", val_loader)]:
        correct = 0;
        total = 0
        with torch.no_grad():
            for imgs, labels in loader:
                imgs = imgs.to(device)
                labels = labels.to(device)
                outputs = model(imgs)
                _, pred = torch.max(outputs, dim=1)
                total += labels.shape[0]
                correct += int((pred == labels).sum())
        print("Accuracy {}: {:.2f}".format(name, correct / total))
        accdict[name] = correct / total
    return accdict
